import os
import shutil
import subprocess
import sys
import argparse
import platform

# 定义颜色变量
RED = "\033[31m"
GREEN = "\033[32m"
YELLOW = "\033[33m"
BLUE = "\033[34m"
PLAIN = "\033[0m"

# 定义软件源列表
MIRROR_LIST = {
    "阿里云": "mirrors.aliyun.com",
    "腾讯云": "mirrors.tencent.com",
    "华为云": "repo.huaweicloud.com",
    "网易": "mirrors.163.com",
    "火山引擎": "mirrors.volces.com",
    "清华大学": "mirrors.tuna.tsinghua.edu.cn",
    "北京大学": "mirrors.pku.edu.cn",
    "浙江大学": "mirrors.zju.edu.cn",
    "南京大学": "mirrors.nju.edu.cn",
    "兰州大学": "mirror.lzu.edu.cn",
    "上海交通大学": "mirror.sjtu.edu.cn",
    "重庆邮电大学": "mirrors.cqupt.edu.cn",
    "中国科学技术大学": "mirrors.ustc.edu.cn",
    "中国科学院软件研究所": "mirror.iscas.ac.cn",
}

# 定义系统类型
SYSTEM_DEBIAN = "Debian"
SYSTEM_UBUNTU = "Ubuntu"
SYSTEM_REDHAT = "RedHat"
SYSTEM_CENTOS = "CentOS"
SYSTEM_FEDORA = "Fedora"
SYSTEM_OPENSUSE = "openSUSE"
SYSTEM_ARCH = "Arch"
SYSTEM_GENTOO = "Gentoo"

# 定义文件路径
FILE_DEBIAN_SOURCE_LIST = "/etc/apt/sources.list"
FILE_YUM_REPOS = "/etc/yum.repos.d"
FILE_ZYPP_REPOS = "/etc/zypp/repos.d"
FILE_PACMAN_MIRRORLIST = "/etc/pacman.d/mirrorlist"

# 定义函数：运行系统命令
def run_command(command, shell=False):
    try:
        subprocess.run(command, shell=shell, check=True)
        print(f"{GREEN}命令执行成功: {command}{PLAIN}")
    except subprocess.CalledProcessError as e:
        print(f"{RED}命令执行失败: {command}\n错误信息: {e}{PLAIN}")
        sys.exit(1)

# 定义函数：检测系统类型
def detect_system():
    system_info = platform.freedesktop_os_release()
    system_name = system_info.get("NAME", "")
    system_version = system_info.get("VERSION_ID", "")
    return system_name, system_version

# 定义函数：备份原有软件源
def backup_sources(system_name):
    if system_name in [SYSTEM_DEBIAN, SYSTEM_UBUNTU]:
        backup_file(FILE_DEBIAN_SOURCE_LIST)
    elif system_name in [SYSTEM_REDHAT, SYSTEM_CENTOS, SYSTEM_FEDORA]:
        backup_dir(FILE_YUM_REPOS)
    elif system_name == SYSTEM_OPENSUSE:
        backup_dir(FILE_ZYPP_REPOS)
    elif system_name == SYSTEM_ARCH:
        backup_file(FILE_PACMAN_MIRRORLIST)

# 定义函数：备份文件或目录
def backup_file(file_path):
    if os.path.exists(file_path):
        backup_path = f"{file_path}.bak"
        shutil.copy2(file_path, backup_path)
        print(f"{GREEN}已备份文件: {file_path} -> {backup_path}{PLAIN}")

def backup_dir(dir_path):
    if os.path.exists(dir_path):
        backup_path = f"{dir_path}.bak"
        shutil.copytree(dir_path, backup_path)
        print(f"{GREEN}已备份目录: {dir_path} -> {backup_path}{PLAIN}")

# 定义函数：更换软件源
def change_sources(system_name, mirror):
    if system_name in [SYSTEM_DEBIAN, SYSTEM_UBUNTU]:
        change_debian_sources(mirror)
    elif system_name in [SYSTEM_REDHAT, SYSTEM_CENTOS, SYSTEM_FEDORA]:
        change_redhat_sources(mirror)
    elif system_name == SYSTEM_OPENSUSE:
        change_opensuse_sources(mirror)
    elif system_name == SYSTEM_ARCH:
        change_arch_sources(mirror)

# 定义函数：更换 Debian/Ubuntu 软件源
def change_debian_sources(mirror):
    with open(FILE_DEBIAN_SOURCE_LIST, "w") as f:
        f.write(f"deb http://{mirror} {detect_system()[1]} main\n")
    print(f"{GREEN}已更换 Debian/Ubuntu 软件源为: {mirror}{PLAIN}")

# 定义函数：更换 RedHat/CentOS/Fedora 软件源
def change_redhat_sources(mirror):
    for file in os.listdir(FILE_YUM_REPOS):
        if file.endswith(".repo"):
            file_path = os.path.join(FILE_YUM_REPOS, file)
            with open(file_path, "r+") as f:
                content = f.read()
                content = content.replace("baseurl=", f"baseurl=http://{mirror}/")
                f.seek(0)
                f.write(content)
                f.truncate()
    print(f"{GREEN}已更换 RedHat/CentOS/Fedora 软件源为: {mirror}{PLAIN}")

# 定义函数：更换 openSUSE 软件源
def change_opensuse_sources(mirror):
    for file in os.listdir(FILE_ZYPP_REPOS):
        if file.endswith(".repo"):
            file_path = os.path.join(FILE_ZYPP_REPOS, file)
            with open(file_path, "r+") as f:
                content = f.read()
                content = content.replace("baseurl=", f"baseurl=http://{mirror}/")
                f.seek(0)
                f.write(content)
                f.truncate()
    print(f"{GREEN}已更换 openSUSE 软件源为: {mirror}{PLAIN}")

# 定义函数：更换 Arch Linux 软件源
def change_arch_sources(mirror):
    with open(FILE_PACMAN_MIRRORLIST, "w") as f:
        f.write(f"Server = http://{mirror}/$repo/os/$arch\n")
    print(f"{GREEN}已更换 Arch Linux 软件源为: {mirror}{PLAIN}")

# 定义函数：更新软件包
def update_packages(system_name):
    if system_name in [SYSTEM_DEBIAN, SYSTEM_UBUNTU]:
        run_command(["apt-get", "update"], shell=True)
        run_command(["apt-get", "upgrade", "-y"], shell=True)
    elif system_name in [SYSTEM_REDHAT, SYSTEM_CENTOS, SYSTEM_FEDORA]:
        run_command(["dnf", "makecache"], shell=True)
        run_command(["dnf", "upgrade", "-y"], shell=True)
    elif system_name == SYSTEM_OPENSUSE:
        run_command(["zypper", "ref"], shell=True)
        run_command(["zypper", "update", "-y"], shell=True)
    elif system_name == SYSTEM_ARCH:
        run_command(["pacman", "-Sy"], shell=True)
        run_command(["pacman", "-Su", "-y"], shell=True)

# 主函数
def main():
    parser = argparse.ArgumentParser(description="更换 Linux 系统软件源")
    parser.add_argument("--mirror", help="指定软件源地址", required=True)
    args = parser.parse_args()

    system_name, _ = detect_system()
    print(f"{BLUE}检测到系统: {system_name}{PLAIN}")

    backup_sources(system_name)
    change_sources(system_name, args.mirror)
    update_packages(system_name)

    print(f"{GREEN}软件源更换完成！{PLAIN}")

if __name__ == "__main__":
    main()